// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*!
\file
\brief Implementation of SNP-Major Allelic Bitstring array 
*/

#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#include "bitstring.h"
#include "bmapfile.h"
#include "copyfile.h"
#include "locus.h"
#include "spedfile.h"
#include "suffixes.h"
#include "types.h"
#include "utils.h"

int spedfile_create(struct mapfile *map, struct selected_options *ops) {
	FILE *fout = NULL;
	BitString bs = NULL;
	int bsAllocatedStorageInBytes = 0, required_bits = 0, nSNPs = 0, nSamples = 0;
	int i, j;

	if (map == NULL || ops == NULL) return 0;
	else if (map->nSNP == 0) return 0;
	else if (strlen(ops->szSpedFileName) == 0) return 0;

	// Calc required bits based on sample size using the nOne size index.
	// This value just reflects the number of  samples in the input array.
	nSNPs = map->nSNP;
	nSamples = map->SNP[0]->nOne;
	required_bits = nSamples * 2;

	if (required_bits == 0) return 0;

	// Allocate the bit string.
	bs = BitString_create(required_bits, &bsAllocatedStorageInBytes);
	if (bs == NULL) return 0;

	fout = fopen(ops->szSpedFileName, "wb");
	if (fout == NULL) {
		if (bs != NULL) free(bs);
		return 0;
	}

	// Firstly write 3 'int' values to the file.
	// nSNPs, nSamples & bytes allocated to the bitstring.
	// Keeping the bitstring size incase meet any 64-bit weirdness down the line.
	if (fwrite(&nSNPs, sizeof(int), 1, fout) != 1) {
		if (fout != NULL) fclose(fout);
		if (bs != NULL) free(bs);
		return 0;
	}
	if (fwrite(&nSamples, sizeof(int), 1, fout) != 1) {
		if (fout != NULL) fclose(fout);
		if (bs != NULL) free(bs);
		return 0;
	}
	if (fwrite(&bsAllocatedStorageInBytes, sizeof(int), 1, fout) != 1) {
		if (fout != NULL) fclose(fout);
		if (bs != NULL) free(bs);
		return 0;
	}

	// Loop through the SNP array, writing the genomic data to the bit string.
	for (i = 0; i < map->nSNP; i++) {
		long bitpos = 0;

		// Write the bitstring of genomic data.
		BitStringset0(bs, 0, required_bits);
		for (j =0; j < nSamples; j++) {
			if (map->SNP[i]->one[j]) BitStringset1(bs, bitpos);
			bitpos++;
			if (map->SNP[i]->two[j]) BitStringset1(bs, bitpos);
			bitpos++;
		}

		if (fwrite(bs,bsAllocatedStorageInBytes,1,fout) != 1) {
			if (fout != NULL) fclose(fout);
			if (bs != NULL) free(bs);
			return 0;
		}
	} 

	if (fout != NULL) fclose(fout);
	if (bs != NULL) free(bs);

	return 1;
}

char* spedfile_load(struct selected_options *ops, int *nSNPs, int *nSamples, int *nRecordBytes) {
	char *bytes = NULL;
	FILE *in = NULL;
	long nRequiredBytes = 0;

	if (ops == NULL || nSamples == NULL || nSNPs == NULL || nRecordBytes == NULL) 
		return NULL;
	else if (strlen(ops->szSpedFileName) == 0) 
		return NULL;

	in = fopen(ops->szSpedFileName, "rb");
	if (in == NULL) return NULL;

	// No. SNPs
	if (fread(nSNPs, sizeof(int), 1, in) != 1) {
		if (in != NULL) fclose(in);
		return NULL;
	}

	// Sample size
	if (fread(nSamples, sizeof(int), 1, in) != 1) {
		if (in != NULL) fclose(in);
		return NULL;
	}

	// Record size in bytes
	if (fread(nRecordBytes, sizeof(int), 1, in) != 1) {
		if (in != NULL) fclose(in);
		return NULL;
	}

	// Load the remaining IPED file into memory.
	nRequiredBytes = (*nSNPs) * (*nRecordBytes);
	bytes = (char*) malloc(nRequiredBytes * sizeof(char));
	if (bytes != NULL) {
		// Failed to load the buffer into memory so leave deallocate
		// and quit.
		if (fread(bytes, nRequiredBytes, 1, in) != 1) {
			free(bytes);	
			return NULL;
		}
	}
	if (in != NULL) fclose(in);

	return bytes;
}

// Locally temporarily appending rank to the run ID to increase the uniqueness of the output file name.
int spedfile_create_and_copy(struct selected_options *ops, struct mapfile *m, int rank) {
	BOOL jobDone = FALSE;
	char tmp_szRunId[TOKEN_LENGTH], remoteFilepath[FILE_PATH_LENGTH];

	if (m == NULL || ops == NULL) return 0;

	strcpy(tmp_szRunId, ops->szRunId);
	sprintf(ops->szRunId, "%s_%i", tmp_szRunId, rank);

	sprintf(ops->szSpedFileName, "%s%s%s_%i", ops->ldir, ops->szRunId, SUFFIX_SPED_FILE, rank);
	sprintf(remoteFilepath, "%s%s%s", ops->rdir, tmp_szRunId, SUFFIX_SPED_FILE);

	if (spedfile_create(m, ops))  if (copyAndDeleteFile(ops->szSpedFileName, remoteFilepath)) jobDone = TRUE;

	strcpy(ops->szRunId, tmp_szRunId); // Restore the original Run ID

	if (jobDone) return 1;
	else return 0;
}

char* _spedfile_load(struct selected_options *ops, int *nSNPs, int *nSamples, int *nRecordBytes) {
	if (ops == NULL || nSNPs == NULL || nSamples == NULL || nRecordBytes == NULL) return NULL;

	sprintf(ops->szSpedFileName, "%s%s%s", ops->rdir, ops->szRunId, SUFFIX_SPED_FILE);
	return spedfile_load(ops, nSNPs, nSamples, nRecordBytes);
}

int spedfile_to_ipedfile(struct selected_options *ops) {
	BOOL jobDone = FALSE;
	char *spedData = NULL, *ipedData = NULL, localFilepath[FILE_PATH_LENGTH], remoteFilepath[FILE_PATH_LENGTH];
	BOOL *person_record = NULL;
	int nSNPs = 0, nSamples = 0, nSpedRecordSizeInBytes = 0, nIpedRecordSizeInBytes = 0, i, j;
	BitString bs_i = NULL, bs_s = NULL;

	if (ops == NULL) return 0;
	else if (strlen(ops->szRunId) == 0 || strlen(ops->ldir) == 0 || strlen(ops->rdir) == 0 || strlen(ops->szSpedFileName) == 0) return 0;

	spedData = spedfile_load(ops, &nSNPs, &nSamples, &nSpedRecordSizeInBytes);
	if (spedData != NULL && nSNPs > 0 && nSamples > 0 && nSpedRecordSizeInBytes > 0) {
		bs_i = BitString_create(nSNPs * 2, &nIpedRecordSizeInBytes);
		if (bs_i != NULL) free(bs_i);
		if (nIpedRecordSizeInBytes > 0) {
			ipedData = (char*) calloc(nSamples * nIpedRecordSizeInBytes, 1);
			person_record = create_bool_array(nSNPs * 2, FALSE);
			if (ipedData != NULL && person_record != NULL) {
				FILE *out = NULL;
				long person_bitpos = 0;
				for (i = 0; i < nSamples; i++, person_bitpos += 2) {
					bs_i = (BitString) (ipedData + (i * nIpedRecordSizeInBytes));
					for (j = 0; j < nSNPs; j++) {
						long snp_bitpos = j * 2;
						bs_s = (BitString) (spedData + (nSpedRecordSizeInBytes * j));
						person_record[snp_bitpos] = BitStringgetbit(bs_s, person_bitpos);
						person_record[snp_bitpos + 1] = BitStringgetbit(bs_s, person_bitpos + 1);
						if (person_record[snp_bitpos]) BitStringset1(bs_i, snp_bitpos);
						if (person_record[snp_bitpos + 1]) BitStringset1(bs_i, snp_bitpos + 1);
					}
				}

				sprintf(localFilepath, "%s%s%s", ops->ldir, ops->szRunId, SUFFIX_IPED_FILE);
				sprintf(remoteFilepath, "%s%s%s", ops->rdir, ops->szRunId, SUFFIX_IPED_FILE);
				out = fopen(localFilepath, "wb");
				if (out != NULL) {
					if (fwrite(&nSamples, sizeof(nSamples), 1, out) == 1) {
						if (fwrite(&nSNPs, sizeof(nSNPs), 1, out) == 1) {
							if (fwrite(&nIpedRecordSizeInBytes, sizeof(nIpedRecordSizeInBytes), 1, out) == 1) {
								size_t expected_amount = nSamples * nIpedRecordSizeInBytes;
								if (fwrite(ipedData, expected_amount, 1, out) == 1) jobDone = TRUE;
							}
						}
					}
					fclose(out);

					if (jobDone) jobDone = copyFile(localFilepath, remoteFilepath) ? TRUE : FALSE;
					else deleteFile(localFilepath);
				}
			}
		}
	}
	if (jobDone) deleteFile(ops->szSpedFileName); // Not need as data now encoded in IPED form.

	if (spedData != NULL) free(spedData);
	if (ipedData != NULL) free(ipedData);
	if (person_record != NULL) free(person_record);

	if (jobDone) return 1;
	else return 0;
}